import pytest
from app import create_app, db
from app.models.knowledge_base import KnowledgeBase
from app.models.system_prompt import SystemPromptTemplate
from app.utils.temp_upload_dir import temp_upload_dir

@pytest.fixture
def app():
    app = create_app()
    app.config['TESTING'] = True
    app.config['SQLALCHEMY_DATABASE_URI'] = 'postgresql://localhost/test_knowledge_base'
    
    with app.app_context():
        db.create_all()
        yield app
        db.session.remove()
        db.drop_all()

@pytest.fixture
def client(app):
    return app.test_client()

@pytest.fixture
def knowledge_base(app):
    kb = KnowledgeBase(
        name='Test KB',
        description='Test Description',
        embedding_model='test-model'
    )
    db.session.add(kb)
    db.session.commit()
    return kb 

@pytest.fixture
def upload_dir(app):
    with temp_upload_dir(app) as upload_dir:
        yield upload_dir

@pytest.fixture
def mock_ollama(monkeypatch):
    def mock_generate_embedding(*args, **kwargs):
        return [0.1] * 384  # 模拟向量

    monkeypatch.setattr(
        'app.services.embedding_service.EmbeddingService.generate_embedding',
        mock_generate_embedding
    )

@pytest.fixture
def mock_celery(monkeypatch):
    class MockAsyncResult:
        def __init__(self, id):
            self.id = id

    def mock_delay(*args, **kwargs):
        return MockAsyncResult('test-task-id')

    monkeypatch.setattr(
        'app.tasks.document_tasks.process_document.delay',
        mock_delay
    ) 